from fastapi import APIRouter, HTTPException, BackgroundTasks
from typing import Union, Dict, Any
from app.services.medical_share_service import MedicalShareService
from app.tools.redis_client import redis_client
from app.schemas.medical_share_schema import *
from app.utils.logger import logger

router = APIRouter()
service = MedicalShareService()

@router.post("/create", response_model=AlgorithmMiddleResponse)
async def create_task(request: AlgorithmRequest) -> AlgorithmMiddleResponse:
    """创建新的医疗数据共享优化任务"""
    try:
        task_response = AlgorithmMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.model_dump())
        return task_response
    except Exception as e:
        logger.error(f"创建医疗数据共享任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

async def add_task_to_background(
    background_tasks: BackgroundTasks,
    input_params: Dict[str, Any],
    task_id: str
) -> None:
    """添加任务到后台执行队列"""
    background_tasks.add_task(run_task_in_background, input_params, task_id)

@router.post("/execute/{task_id}", response_model=AlgorithmMiddleResponse)
async def execute_task(
    task_id: str,
    background_tasks: BackgroundTasks
) -> AlgorithmMiddleResponse:
    """执行医疗数据共享优化任务"""
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="任务不存在")
            
        input_params = task_data.get("input_params")
        logger.info(f"执行任务的输入参数: {input_params}")
        
        if not input_params:
            raise HTTPException(status_code=400, detail="缺少input_params")
            
        if "medical_share_params" not in input_params:
            raise HTTPException(status_code=400, detail="缺少medical_share_params")
        
        result_key = f'{task_id}_result'
        
        running_status = {
            "task_id": task_id,
            "task_callback_url": task_data.get("task_callback_url"),
            "task_status": TaskStatus.RUNNING,
            "task_progress": 0,
            "input_params": input_params,
            "metrics": []
        }
        await redis_client.set_data(result_key, running_status)
        
        await add_task_to_background(background_tasks, input_params, task_id)
        return AlgorithmMiddleResponse(**running_status)
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"执行医疗数据共享任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/result/{task_id}", response_model=Union[AlgorithmResponse, AlgorithmMiddleResponse])
async def get_result(task_id: str) -> Union[AlgorithmResponse, AlgorithmMiddleResponse]:
    """获取任务执行结果"""
    try:
        result_key = f'{task_id}_result'
        result_data = await redis_client.get_data(result_key)
        if not result_data:
            raise HTTPException(status_code=404, detail="结果不存在")
            
        return (AlgorithmResponse(**result_data) 
                if "output_params" in result_data 
                else AlgorithmMiddleResponse(**result_data))
                
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取医疗数据共享结果失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

async def run_task_in_background(input_params: Dict[str, Any], task_id: str) -> None:
    """后台执行任务"""
    try:
        result = await service.run_medical_share(
            input_params=input_params,
            task_id=task_id
        )
        await redis_client.update_data(f'{task_id}_result', result.model_dump())
    except Exception as e:
        logger.error(f"医疗数据共享后台任务执行失败: {str(e)}")
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e)
        }
        await redis_client.update_data(f'{task_id}_result', error_status) 